#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/1/17 14:05
# @Author  : 杨义乐
# @File    : json_to_txt.py
# @Description : TODO

import json
import os

def json2yolo(json_path, output_dir, category_list):
    with open(json_path, 'r') as f:
        labelme_data = json.load(f)  # 读取JSON数据

    image_width = labelme_data.get('imageWidth', 1)  # 避免为0的情况
    image_height = labelme_data.get('imageHeight', 1)
    l = []

    # 遍历所有的标注形状
    for shape in labelme_data['shapes']:
        label = shape['label']
        if label not in category_list:
            print(f"Skipping shape with label '{label}' not in category list.")
            continue
        category = category_list.index(label)
        points = shape['points']
        if shape['shape_type'] == 'rectangle':
            (x1, y1), (x2, y2) = points
        elif shape['shape_type'] == 'polygon':
            x1, y1 = min(point[0] for point in points), min(point[1] for point in points)
            x2, y2 = max(point[0] for point in points), max(point[1] for point in points)
        else:
            continue
        # 计算YOLO格式所需的中心点和宽高
        x_center = max(0, (x1 + x2) / 2.0 / image_width)
        y_center = max(0, (y1 + y2) / 2.0 / image_height)
        width = max(0, (x2 - x1) / image_width)
        height = max(0, (y2 - y1) / image_height)
        l.append(f"{category} {x_center} {y_center} {width} {height}")

    # 输出txt文件
    output_file = os.path.join(output_dir, os.path.splitext(os.path.basename(json_path))[0] + '.txt')
    with open(output_file, 'w') as f:
        f.write('\n'.join(l))

def process_folder(json_folder, output_folder, category_list):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    for filename in os.listdir(json_folder):
        if filename.endswith('.json'):
            json_path = os.path.join(json_folder, filename)
            json2yolo(json_path, output_folder, category_list)
        print("Conversion completed!")

# 定义类别列表
category_list = ['pip1']

# 输入文件夹和输出文件夹
json_folder = r'C:\Users\10067\PycharmProjects\pipeline\out'  # 修改为输入json文件夹路径
output_txt_folder = r'C:\Users\10067\PycharmProjects\pipeline\label'  # 修改为输出txt文件夹路径

# 处理文件夹中的所有json文件
process_folder(json_folder, output_txt_folder, category_list)